package com.test.chatrobot3;

import org.deeplearning4j.datasets.iterator.impl.ListDataSetIterator;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Nesterovs;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.io.File;
import java.util.*;

public class SimpleChatBot {

    // 词汇表和回答映射
    private static List<String> vocabulary = Arrays.asList("hello", "what", "your", "name", "how", "are", "you");
    private static Map<Integer, String> responseMap = new HashMap<>();
    static {
        responseMap.put(0, "Hello! How can I help you?");
        responseMap.put(1, "My name is ChatBot.");
        responseMap.put(2, "I'm fine, thank you!");
    }

    public static void main(String[] args) throws Exception {
        // 训练模型
        MultiLayerNetwork model = trainModel();
        // 保存模型
        ModelSerializer.writeModel(model, new File("chatbot-model.zip"), true);

        // 加载模型进行预测
        MultiLayerNetwork loadedModel = ModelSerializer.restoreMultiLayerNetwork("chatbot-model.zip");

        // 与用户交互
        Scanner scanner = new Scanner(System.in);
        while (true) {
            System.out.print("You: ");
            String input = scanner.nextLine();
            if ("exit".equalsIgnoreCase(input)) break;

            INDArray features = textToFeatureVector(input);
            INDArray output = loadedModel.output(features);
            int predictedClass = Nd4j.argMax(output, 1).getInt(0);
            System.out.println("Bot: " + responseMap.get(predictedClass));
        }
        scanner.close();
    }

    private static MultiLayerNetwork trainModel() {
        // 创建训练数据
        List<DataSet> trainingData = new ArrayList<>();
        // 样本数据（输入问题，输出类别）
        addExample(trainingData, "hello", 0);
        addExample(trainingData, "what your name", 1);
        addExample(trainingData, "how are you", 2);

        DataSetIterator iterator = new ListDataSetIterator<>(trainingData, 3);

        // 配置神经网络
        int numInput = vocabulary.size();
        int numOutput = responseMap.size();

        MultiLayerConfiguration config = new NeuralNetConfiguration.Builder()
                .seed(1234)
                .updater(new Nesterovs(0.1, 0.9))
                .list()
                .layer(new DenseLayer.Builder()
                        .nIn(numInput)
                        .nOut(5)
                        .activation(Activation.RELU)
                        .weightInit(WeightInit.XAVIER)
                        .build())
                .layer(new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD)
                        .nIn(5)
                        .nOut(numOutput)
                        .activation(Activation.SOFTMAX)
                        .weightInit(WeightInit.XAVIER)
                        .build())
                .build();

        MultiLayerNetwork model = new MultiLayerNetwork(config);
        model.init();
        model.fit(iterator, 100); // 训练100轮

        return model;
    }

    private static void addExample(List<DataSet> list, String text, int label) {
        INDArray features = textToFeatureVector(text);
        INDArray labels = Nd4j.zeros(1, responseMap.size());
        labels.putScalar(label, 1);
        list.add(new DataSet(features, labels));
    }

    private static INDArray textToFeatureVector(String text) {
        INDArray vector = Nd4j.zeros(1, vocabulary.size());
        String[] tokens = text.toLowerCase().split(" ");
        for (String token : tokens) {
            int index = vocabulary.indexOf(token);
            if (index != -1) {
                vector.putScalar(index, 1.0);
            }
        }
        return vector;
    }
}